#include"hfsolver_ho.h"
void SHFSolver_HO::init()
{
  //initialize HF Orbitals with hamiltonian basis
  int totalOrbitals=pSystem_Tz->Orbitals.size();
  HFOrbitals.resize(totalOrbitals);
  for(int i=0;i<pSystem_Tz->Groups.size();i++)
    {
      int numOrbInAGroup=pSystem_Tz->Groups[i].size();
      for(int j=0;j<numOrbInAGroup;j++)
	{
	  int orb=pSystem_Tz->Groups[i][j];
	  const OrbitalType & Orb=pSystem_Tz->Orbitals[orb];
	  HFOrbitals[orb].set(i,Orb.l,Orb.jj,Orb.tz,Orb.e);
	  HFOrbitals[orb].Coeff.resize(numOrbInAGroup);
	  //clear all coeff
	  for(int indice=0;indice<numOrbInAGroup;indice++)
	    {
	      HFOrbitals[orb].Coeff[indice]=0;
	    }
	  //set initial coeff
	  HFOrbitals[orb].Coeff[j]=DataType(1);
	}
    }
  fillOrbitals();
}

SHFSolver_HO::DataType SHFSolver_HO::SHFDiagram(int bra,int ket) const
{
  int jj=pSystem_Tz->Orbitals[bra].jj;
  DataType val2B=0;
  for(int i=0;i<HFOrbitals.size();i++)
    {
      if(HFOrbitals[i].state) continue;//skip particles
      int Jmin=abs(jj - HFOrbitals[i].jj)/2;
      int Jmax=(jj + HFOrbitals[i].jj)/2;


      int GroupIndex=HFOrbitals[i].GroupIndex;
      int dim=pSystem_Tz->Groups[GroupIndex].size();

      for(int k=0;k<dim;k++)
	for(int l=0;l<dim;l++)
	  {
	    for(int J=Jmin;J<=Jmax;J++)
	      {
		int brak=pSystem_Tz->Groups[GroupIndex][k];
		int ketl=pSystem_Tz->Groups[GroupIndex][l];
		val2B += getTwoBody(bra,brak,ket,ketl,J) * HFOrbitals[i].Coeff[k] * HFOrbitals[i].Coeff[l] * HFOrbitals[i].OccuP * DataType(2*J+1);
	      }
	  }
    }
  val2B/=(jj+1);
  return val2B;
}

SHFSolver_HO::DataType SHFSolver_HO::getVHF(int bra,int ket) const
{
  int jj=pSystem_Tz->Orbitals[bra].jj;
  DataType val2B=0;
  for(int i=0;i<HFOrbitals.size();i++)
    {
      if(HFOrbitals[i].state) continue;//skip particles
      int Jmin=abs(jj - HFOrbitals[i].jj)/2;
      int Jmax=(jj + HFOrbitals[i].jj)/2;


      int GroupIndex=HFOrbitals[i].GroupIndex;
      int dim=pSystem_Tz->Groups[GroupIndex].size();

      for(int k=0;k<dim;k++)
	for(int l=0;l<dim;l++)
	  {
	    for(int J=Jmin;J<=Jmax;J++)
	      {
		int brak=pSystem_Tz->Groups[GroupIndex][k];
		int ketl=pSystem_Tz->Groups[GroupIndex][l];
		val2B += getV(bra,brak,ket,ketl,J) * HFOrbitals[i].Coeff[k] * HFOrbitals[i].Coeff[l] * HFOrbitals[i].OccuP * DataType(2*J+1);
	      }
	  }
    }
  val2B/=(jj+1);
  return val2B;
}


void SHFSolver_HO::fillOrbitals()
{
  stable_sort( HFOrbitals.begin(),HFOrbitals.end() );
  int resZ=Z;
  int resN=A-Z;
  for(int i=0;i<HFOrbitals.size();i++)
    {
      //fill proton
      if(HFOrbitals[i].tz==-1)
	{
	  int num=HFOrbitals[i].jj+1;
	  resZ-=num;
	  if(resZ >= 0)
	    {
	      HFOrbitals[i].state=0;//hole
	      HFOrbitals[i].OccuP=1.0;
	    }
	  else if( (resZ + num) > 0)
	    {
	      HFOrbitals[i].state=0;//hole, but not fully occupied.
	      HFOrbitals[i].OccuP=double(resZ + num)/num;
	    }
	  else
	    {
	      HFOrbitals[i].state=1;//particle
	      HFOrbitals[i].OccuP=0.0;
	    }
	}
      //fill neutron
      if(HFOrbitals[i].tz==1)
	{
	  int num=HFOrbitals[i].jj+1;
	  resN-=num;
	  if(resN >= 0)
	    {
	      HFOrbitals[i].state=0;//hole
	      HFOrbitals[i].OccuP=1.0;
	    }
	  else if( (resN + num) > 0)
	    {
	      HFOrbitals[i].state=0;//hole, but not fully occupied.
	      HFOrbitals[i].OccuP=double(resN + num)/num;
	    }
	  else
	    {
	      HFOrbitals[i].state=1;//particle
	      HFOrbitals[i].OccuP=0.0;
	    }

	}
    }
}

void SHFSolver_HO::diag()
{
  int totalOrbitals=HFOrbitals.size();
  vector< SHF_Orbital<DataType> > newHFOrbitals(totalOrbitals);
  
  int num=0;
  int totalGroups=pSystem_Tz->Groups.size();
  for(int i=0;i<totalGroups;i++)
    {
      int orb=pSystem_Tz->Groups[i][0];
      const OrbitalType & Orb=pSystem_Tz->Orbitals[orb];
      int l=Orb.l;
      int jj=Orb.jj;
      int tz=Orb.tz;
      
      int dim=pSystem_Tz->Groups[i].size();
      Mat mat(dim,dim);
      for(int a=0;a<dim;a++)
	{
	  for(int b=a;b<dim;b++)
	    {
	      mat(a,b)=HFhamil(i,a,b);
	      mat(b,a)=mat(a,b);
	    }
	}
      SelfAdjointEigenSolver<Mat> es(mat);
      Vec eigenVals=es.eigenvalues();
      Mat eigenVecs=es.eigenvectors();
      for(int j=0;j<dim;j++)
	{
	  newHFOrbitals[num+j].set(i,l,jj,tz,eigenVals[j]);
	  newHFOrbitals[num+j].Coeff.resize(dim);
	  //Coeff
	  for(int k=0;k<dim;k++)
	    newHFOrbitals[num+j].Coeff[k]=eigenVecs(k,j);
	}
      num+=dim;
    }
  HFOrbitals=newHFOrbitals;
  fillOrbitals();
}

void SHFSolver_HO::iter()
{
  double sigma=0;
  double eps=1e-6;
  const int maxIter=200;
  init();
  int totalOrbitals=HFOrbitals.size();
  int iter=0;
  //  cout<<"begin HF. iteration...\n";
  do
    {     
      vector< SHF_Orbital<DataType> > lastHFOrbitals=HFOrbitals;
      //printOrbitals(6);
      diag();
      sigma=0;
      for(int i=0;i<totalOrbitals;i++)
  	{
  	  sigma+=abs(lastHFOrbitals[i].e-HFOrbitals[i].e);
  	}
      sigma/=totalOrbitals;
      ++iter;
      //      cout<<iter<<"\t"<<sigma<<endl;
    }
  while(sigma>eps && iter<maxIter);
  if(iter==maxIter)
    std::cout<<"sigma:"<<sigma<<", convergence failed!\n";
  //  cout<<"end  HF. iteration\n";
}



void SHFSolver_HO::calEnergy()
{
  Energy=0;
  for(int i=0;i<HFOrbitals.size();i++)
    {
      if(HFOrbitals[i].state) continue;//skip particles
      int dim=HFOrbitals[i].Coeff.size();
      int group=HFOrbitals[i].GroupIndex;
      double Possi=HFOrbitals[i].OccuP;
      int jj1=HFOrbitals[i].jj+1;
      for(int a=0;a<dim;a++)
	{
	  for(int b=0;b<dim;b++)
	    {
	      int bra=pSystem_Tz->Groups[group][a];
	      int ket=pSystem_Tz->Groups[group][b];	      
	      Energy+=( getOneBody(bra,ket)+0.5*SHFDiagram(bra,ket) )*Possi*HFOrbitals[i].Coeff[a]*HFOrbitals[i].Coeff[b]*DataType(jj1);
	    }
	}
    }
  Energy-=beta*pSystem_Tz->hbar_omega;
}

SHFSolver_HO::DataType SHFSolver_HO::hf_r(int i,double r) const
{
  double hfr=0;
  int GroupIndex=HFOrbitals[i].GroupIndex;
  int dim=pSystem_Tz->Groups[GroupIndex].size();
  int l=HFOrbitals[i].l;
  double b=pSystem_Tz->lengthb;
  for(int j=0;j<dim;j++)
    {
      int orb=pSystem_Tz->Groups[GroupIndex][j];
      int n=pSystem_Tz->Orbitals[orb].n;
      hfr +=  HFOrbitals[i].Coeff[j] * ho_r(n,l,b,r);
    }
  return hfr;
}


//*********************************************************************

